from pymilvus import connections, Collection, utility
from app.core.config import settings
import logging

logger = logging.getLogger(__name__)

class MilvusHandler:
    _connected = False

    @classmethod
    def connect(cls):
        if not cls._connected:
            try:
                connections.connect(
                    alias="default",
                    host=settings.MILVUS_HOST,
                    port=settings.MILVUS_PORT,
                    user=settings.MILVUS_USER,
                    password=settings.MILVUS_PASSWORD
                )
                cls._connected = True
                logger.info(f"Connected to Milvus at {settings.MILVUS_HOST}:{settings.MILVUS_PORT}")
            except Exception as e:
                logger.error(f"Failed to connect to Milvus: {e}")
                raise

    @classmethod
    def get_collection(cls, collection_name: str = None) -> Collection:
        """
        获取 Milvus Collection 对象。
        确保在使用前调用过 connect()，或者在这里做懒加载连接。
        """
        cls.connect()
        name = collection_name or settings.MILVUS_COLLECTION_NAME
        
        if not utility.has_collection(name):
            raise ValueError(f"Collection '{name}' does not exist. Please initialize it first.")
            
        return Collection(name)

# 便捷使用的函数
def get_milvus_client():
    """
    对于 pymilvus 2.3+，也可以使用高层级的 MilvusClient 类
    这里返回连接后的 Collection 获取器或连接对象，视具体使用习惯而定
    """
    MilvusHandler.connect()
    return connections